#ifndef TINYNDARRAY_COMMON_H
#define TINYNDARRAY_COMMON_H

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <half.hpp>
#include <initializer_list>
#include <iostream>
#include <list>
#include <memory>
#include <mutex>
#include <random>
#include <stdexcept>
#include <thread>
#include <utility>
#include <variant>
#include <vector>
using float32_t = float;
using float64_t = double;
using float16_t = half_float::half;
namespace tinyndarray {
// 辅助类型定义
using InitShape = std::initializer_list<int>;
using Shape = std::vector<int>;
using Index = std::vector<int>;
using SliceIndex = std::vector<std::pair<int, int>>;
using Axis = std::vector<int>;
using IList = std::initializer_list<int>;
using QuantInfo = std::pair<float, int>;
// 数据类型枚举
enum DtypeTag {
    INT8,
    UINT8,
    INT16,
    UINT16,
    INT32,
    UINT32,
    INT64,
    UINT64,
    FP16,
    FP32,
    FP64,
    BOOL,
    QINT8,
    QUINT8,
    UNKNOWN
};

bool is_quant_type(DtypeTag tag);

template <typename T>
T quantize_linear(float x, QuantInfo quant_meta, T qmin, T qmax) {
    auto scale = quant_meta.first;
    auto zp = quant_meta.second;
    auto q = std::round((x - zp) / scale);
    if (q > qmax) q = qmax;
    if (q < qmin) q = qmin;
    return static_cast<T>(q);
}

template <typename T>
T quantize_linear(float x, QuantInfo quant_meta) {
    auto qmin = std::numeric_limits<T>::min();
    auto qmax = std::numeric_limits<T>::max();
    return quantize_linear(x, quant_meta, qmin, qmax);
}

template <typename T>
float dequantize_linear(T q, QuantInfo quant_meta) {
    auto scale = quant_meta.first;
    auto zp = quant_meta.second;
    auto x = q * scale + zp;
    return x;
}

struct Dtype {
    Dtype() : tag(UNKNOWN) {}
    Dtype(DtypeTag _tag) : tag(_tag) {}
    Dtype(DtypeTag _tag, QuantInfo quant_info) : tag(_tag), meta_info(quant_info) {}
    DtypeTag tag;
    std::variant<QuantInfo> meta_info;
    QuantInfo get_quant_info() {
        return std::get<QuantInfo>(meta_info);
    }
    bool operator==(const Dtype& other) {
        return tag == other.tag;
    }
    bool operator!=(const Dtype& other) {
        return tag != other.tag;
    }
};

size_t get_type_size(Dtype dtype);

template <typename T>
inline Dtype map_ctype_to_datatype() {
    if (std::is_same_v<T, float>) {
        return FP32;
    } else if (std::is_same_v<T, double>) {
        return FP64;
    } else if (std::is_same_v<T, int>) {
        return INT32;
    }
    throw std::runtime_error("Unsupported datatype");
    return UNKNOWN;
}

size_t calculate_offset(const Index& index, const Shape& shape, const std::vector<size_t>& strides);
std::vector<size_t> compute_strides(const Shape& shape, size_t element_size);

void GetParallelParams(int size, int& n_workers, int& n_batch, int& batch_size);

std::vector<int> ComputeChildSizes(const Shape& shape);

bool operator==(const Shape& lshape, const Shape& rshape);
bool operator!=(const Shape& lshape, const Shape& rshape);

template <typename T>
std::string vector_int_to_string(const std::vector<T>& int_vec);

extern template std::string vector_int_to_string(const std::vector<int>& int_vec);

bool is_index_in_range(const Index& index, const Shape& shape);
}  // namespace tinyndarray
#endif